算子优化介绍

// 这就是一个 Kernel —— 使用 __global__ 关键字声明
__global__ void add(float *a, float *b, float *c, int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) c[i] = a[i] + b[i];
}

// CPU 调用 Kernel，在 GPU 上并行执行
add<<>>(a, b, c, n);

__global__ void example(float *input, float *output) {
    // 1. 声明共享内存（Block 内线程共享）
    __shared__ float tile[256];
    
    // 2. 从显存读取到共享内存
    tile[threadIdx.x] = input[blockIdx.x * 256 + threadIdx.x];
    __syncthreads();  // 同步，确保数据加载完成
    
    // 3. 从共享内存读到寄存器，进行计算
    float val = tile[threadIdx.x];  // val 在寄存器中
    val = val * 2.0f + 1.0f;        // 计算在寄存器中完成
    
    // 4. 结果写回显存
    output[blockIdx.x * 256 + threadIdx.x] = val;
}

# LayerNorm 中的规约
mean = x.mean(dim=-1)      # 规约：求均值
var = x.var(dim=-1)        # 规约：求方差

# Softmax 中的规约
max_val = x.max(dim=-1)    # 规约：求最大值（数值稳定性）
sum_exp = exp(x - max_val).sum(dim=-1)  # 规约：求和

# Loss 计算
loss = losses.mean()       # 规约：batch 内求均值

# Attention
attn_weights = softmax(scores, dim=-1)  # 内含 max 和 sum 规约

__shared__ float sdata[256];  // 共享内存，用于线程间通信

// 1. 每个线程加载自己的数据到共享内存
sdata[tid] = input[i];
__syncthreads();  // 同步，确保所有线程加载完成

// 2. 并行规约 - 树形结构，每轮活跃线程减半
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
    if (tid < s) {
        sdata[tid] += sdata[tid + s];  // 读取其他线程的数据并累加
    }
    __syncthreads();  // 每轮都要同步
}

// 3. 线程0持有最终结果，写回显存
if (tid == 0) output[blockIdx.x] = sdata[0];

// ❌ 未融合：两个独立 Kernel，两次显存读写
__global__ void add(float *a, float *b, float *c, int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) c[i] = a[i] + b[i];
}

__global__ void relu(float *x, int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) x[i] = x[i] > 0 ? x[i] : 0;
}

// 调用两次 Kernel
add<<>>(a, b, c, n);  // 写 c 到显存
relu<<>>(c, n);       // 读 c，再写回

// ✅ 融合后：一个 Kernel，一次显存写入
__global__ void add_relu(float *a, float *b, float *c, int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        float sum = a[i] + b[i];       // sum 在寄存器中
        c[i] = sum > 0 ? sum : 0;     // 直接计算后写出
    }
}

// 只调用一次 Kernel
add_relu<<>>(a, b, c, n);

# PyTorch 原始写法（3个独立算子）
x = conv(input)      # 卷积: 读 input, 写 x
x = batch_norm(x)    # 归一化: 读 x, 写 x
x = relu(x)          # 激活: 读 x, 写 x

# 融合后（1个融合算子）
x = conv_bn_relu(input)  # 读 input, 写 x
                         # 中间结果在寄存器中传递